#include "Avalon/Services/ServiceProtocol.hpp"
#include "Avalon/Base/Algorithm.hpp"
#include "Avalon/IO/ClientConnection.hpp"
#include "Avalon/IO/Reader.hpp"
#include "Avalon/IO/Writer.hpp"
#include "Avalon/Services/Service.hpp"

using namespace Avalon;
using namespace Avalon::IO;
using namespace Avalon::Serialization;
using namespace Avalon::Services;
using namespace Avalon::Threading;
using namespace boost;
using namespace std;

ServiceProtocol::ServiceProtocol() {}

ServiceProtocol::~ServiceProtocol() {
  Stop();
}

void ServiceProtocol::Initialize(Channel* channel, Serializer* serializer,
    Deserializer* deserializer, map<const type_info*, function_base*>* slots,
    ThreadPool& threadPool) {
  m_channel = channel;
  m_serializer.reset(serializer);
  m_deserializer.reset(deserializer);
  m_slots = slots;
  m_isStarted = false;
  m_isStopping = false;
  m_writeSlot = m_signalHandler.GetSlot<Async<void>::FinishedSlot>(bind(
    &ServiceProtocol::OnWrite, this));
  m_nextRequestId = 1;
  m_signalHandler.Initialize(threadPool);
  m_readResult.SetFinishedSlot(
    m_signalHandler.GetSlot<Async<int>::FinishedSlot>(
    bind(&ServiceProtocol::OnRead, this)), Store(m_readState));
}

Channel* ServiceProtocol::GetChannel() {
  return m_channel;
}

void ServiceProtocol::Start() {
  lock_guard<mutex> lock(m_mutex);
  if(m_isStarted) {
    return;
  }
  m_channel->GetReader().Read(&m_readBuffer, Store(m_readResult));
  m_isStarted = true;
}

void ServiceProtocol::Stop() {
  unique_lock<mutex> lock(m_mutex);
  if(!m_isStarted) {
    return;
  }
  CloseResources(lock);
  if(m_isStarted) {
    m_isStoppedCondition.wait(lock);
  }
}

void ServiceProtocol::SendMessage(const Message& message) {
  lock_guard<mutex> lock(m_mutex);
  m_serializer->Shuttle("message", message);
  WriteToChannel(m_serializer->GetSerializedData());
}

void ServiceProtocol::CloseResources(unique_lock<mutex>& lock) {
  m_isStopping = true;
  if(!m_pendingRequests.empty()) {
    map<int, BaseAsyncHolder*> pendingRequests;
    std::swap(m_pendingRequests, pendingRequests);
    lock.unlock();
    for(map<int, BaseAsyncHolder*>::iterator i = pendingRequests.begin();
        i != pendingRequests.end(); ++i) {
      BaseAsyncHolder* async = NULL;
      std::swap(async, i->second);
      async->SetException(copy_exception(ServiceRequestException(
        "Service protocol has stopped.")));
      delete async;
    }
    lock.lock();
  }
  if(m_readState == Async<int>::PENDING) {
    return;
  }
  if(!m_pendingWrites.empty()) {
    return;
  }
  m_readResult.ResetResult();
  m_isStopping = false;
  m_isStarted = false;
  m_isStoppedCondition.notify_one();
}

void ServiceProtocol::OnRead() {
  unique_lock<mutex> lock(m_mutex);
  m_readState = m_readResult.GetState();
  if(m_readState == Async<int>::EXCEPTION || m_isStopping) {
    m_readResult.ResetResult();
    m_channel->GetConnection().Close();
    CloseResources(lock);
    return;
  }
  int sizeRead = m_readResult.Get();
  m_readResult.ResetResult();
  scoped_ptr<Message> message;
  try {
    m_deserializer->SetDataSource(m_readBuffer.GetData(), sizeRead);
    m_deserializer->Shuttle("message", message);
    if(message == NULL) {
      m_channel->GetConnection().Close();
      CloseResources(lock);
      return;
    }
  } catch(std::exception&) {
    m_channel->GetConnection().Close();
    CloseResources(lock);
    return;
  }
  ProcessMessage(message.get());
  m_readBuffer.Reset();
  m_channel->GetReader().Read(&m_readBuffer, Store(m_readResult));
}

void ServiceProtocol::OnWrite() {
  unique_lock<mutex> lock(m_mutex);
  Async<void>* result = m_pendingWrites.front();
  m_pendingWrites.pop_front();
  Async<void>::State writeState = result->GetState();
  delete result;
  if(m_isStopping && m_pendingWrites.empty()) {
    CloseResources(lock);
    return;
  }
  if(writeState == Async<int>::EXCEPTION) {
    m_channel->GetConnection().Close();
    CloseResources(lock);
  }
}

void ServiceProtocol::ProcessMessage(Message* message) {
  ServiceMessage* serviceMessage = dynamic_cast<ServiceMessage*>(message);
  if(serviceMessage != NULL && serviceMessage->IsResponseMessage()) {
    map<int, BaseAsyncHolder*>::iterator responseIterator =
      m_pendingRequests.find(serviceMessage->GetRequestId());
    if(responseIterator != m_pendingRequests.end()) {
      BaseAsyncHolder* result = responseIterator->second;
      m_pendingRequests.erase(responseIterator);
      serviceMessage->SetAsync(result);
      delete result;
    }
    return;
  }
  if(m_slots == NULL) {
    return;
  }
  map<const type_info*, function_base*>::iterator slotIterator = m_slots->find(
    &typeid(*message));
  if(slotIterator == m_slots->end()) {
    return;
  }
  message->EmitSignal(slotIterator->second, this);
}

void ServiceProtocol::WriteToChannel(const Buffer& buffer) {
  Async<void>::State writeState;
  Async<void>* writeResult = new Async<void>();
  m_channel->GetWriter().Write(buffer, Store(*writeResult));
  writeResult->SetFinishedSlot(m_signalHandler.GetSlot<
    Async<void>::FinishedSlot>(m_writeSlot), Store(writeState));
  if(writeState == Async<void>::PENDING) {
    m_pendingWrites.push_back(writeResult);
  } else {
    delete writeResult;
  }
}
